from airflow.patsnap.util.ding import Ding
from airflow.models.idata_alert import Alert
from airflow.models.idata_task_alert import TaskAlert
from airflow.models.idata_workflow_alert import WorkflowAlert
from airflow.utils.session import provide_session

alert_condition_map = {
    0: '执行失败',
    1: '调度延时',
}

alert_category_map = {
    0: '钉钉告警',
    1: '电话告警',
}


@provide_session
def get_task_alerts(task_id, condition, session=None):
    return session.query(Alert).join(TaskAlert, Alert.id == TaskAlert.alert_id).filter(TaskAlert.task_id == task_id,
                                                                                       TaskAlert.condition == condition).all()


@provide_session
def send_task_alerts(task_id, condition, session=None, msg='', task_name=None):
    try:
        if task_id.find('-'):
            task_id = int(task_id.split('-')[-1])
        for alert in get_task_alerts(task_id, condition, session=session):
            if alert.category == 0:
                text = '### [精卫告警] 任务{} : {} \n {}'.format(alert_condition_map[condition],
                                                       task_name if task_name else task_id,
                                                       msg)
                Ding(alert.content, text).send()
    except Exception as e:
        print(e)


def get_workflow_alerts(workflow_id, condition, session=None):
    alerts = session.query(Alert).join(WorkflowAlert,
                                       Alert.id == WorkflowAlert.alert_id).filter(WorkflowAlert.workflow_id == workflow_id,
                                                                                  WorkflowAlert.condition == condition).all()
    return alerts


def send_workflow_alerts(dag_id, condition, session=None, msg='', workflow_name=None):
    try:
        workflow_id = int(dag_id.split('-')[-1])
        for alert in get_workflow_alerts(workflow_id, condition, session=session):
            if alert.category == 0:
                text = '### [精卫告警] DAG{} : {} \n {}'.format(alert_condition_map[condition], workflow_name, msg)
                Ding(alert.content, text).send()
    except Exception as e:
        print(e)



